# -*-coding=utf-8-*-


import sys
import numpy as np
import random
import time
from ray import Ray
from sphere import Sphere
from wall import Wall
from window import Window
import multiprocessing as mp
from send_email import send_eamil

MAX_DEPTH = 2
RAY_NUMBER = 1
SAMPLE_NUMBER = 1

length = 256
window = Window(length * 2, length * 2)
object_array = [Wall([-length, 0, 0], [1, 0, 0], [0.9, 0, 0], 'diff'),  # left
                Wall([0, -length, 0], [0, 1, 0], [0.99, 0.99, 0.99], 'diff'),  # buttom
                Wall([length, 0, 0], [-1, 0, 0], [0, 0.9, 0.0], 'diff'),  # right
                Wall([0, length, 0], [0, -1, 0], [0.99, 0.99, 0.99], 'diff'),  # top
                Wall([0, 0, -length], [0, 0, 1], [0.99, 0.99, 0.99], 'diff'),  # front
                Wall([0, 0, length], [0, 0, 1], [0.99, 0.99, 0.99], 'back'),  # back
                Sphere([-length/2, -(length/2), -(length/2)], length/2, [0.8, 0.8, 0.8], 'mirr'),  # mirr ball
                Sphere([length/2, -(length/2), -(length/2)], length/2, [0.4, 0.4, 1.0], 'mirr'),  # blue mirr ball
                Sphere([0, length*1.9, -(length/2)], length, [1, 1, 1], 'light')]  # light
light_position = np.array([0, length*0.95, -(length/2)], 'f')
diffuse_part = 0.2

random_vector = np.zeros(3, 'f')


def random_direction(normal):
    """
    生成随机向量，以normal为正方向的半球范围内
    """
    continue_flag = True
    while continue_flag:
        random_vector = [random.uniform(-1, 1), random.uniform(-1, 1), random.uniform(-1, 1)]
        vector_length = np.linalg.norm(random_vector)
        # 向量长度不为0
        if vector_length > 0.00000001:
            continue_flag = False
            random_vector /= vector_length
    if np.dot(random_vector, normal) < 0:
        random_vector = -random_vector
    return random_vector


def intersect(object_array, ray):
    """
    求射线和物体交点
    """
    temp_i = None
    temp_distance = -1.0
    for i, obj in enumerate(object_array):
        info = obj.intersect(ray)
        if info is None:
            continue
        elif info[0] > 0.0:
            # 保存最小正值
            if temp_distance < 0.0 or info[0] < temp_distance:
                temp_i = i
                temp_distance, normal, color, mode = info
    if temp_i is not None:
        temp_position = ray.position + temp_distance * ray.direction
        return (temp_position, normal, color, mode)
    else:
        return None


def light_intersect(ray):
    """
    判断光线能否直接抵达光源
    """
    info = intersect(object_array, ray)
    if info is None:
        return False
    else:
        position, normal, surface_color, mode = info
    if mode == 'light':
        return True
    else:
        return False


def traceray(ray, depth):
    """
    追踪一束光线递归函数
    """
    new_color = np.zeros(3, 'f')
    # 深度过深，返回颜色
    if depth >= MAX_DEPTH:
        return new_color
    info = intersect(object_array, ray)
    # 射线无交点
    if info is None:
        return new_color
    else:
        position, normal, surface_color, mode = info
    # 漫反射，部分光随机选择一个方向反射，剩余参与光源直接照射
    if mode == 'diff':
        for i in range(RAY_NUMBER):
            # 反射部分
            reflect_color = traceray(Ray(position, random_direction(normal), ray.color * surface_color), depth + 1)
            # 光源直接照射部分
            light_ray = Ray(position, light_position - position, ray.color)
            light_flag = light_intersect(light_ray)
            if light_flag:
                from_light_color = ray.color * np.dot(normal, light_ray.direction) * surface_color  # Lambert漫反射模型
                new_color += ((reflect_color * (1-diffuse_part) + from_light_color * diffuse_part) / RAY_NUMBER)
            else:
                new_color += (reflect_color * (1-diffuse_part) / RAY_NUMBER)
        return new_color
    # 镜面反射
    elif mode == 'mirr':
        for i in range(RAY_NUMBER):
            add_color = traceray(Ray(position, ray.direction + 2*normal*(np.abs(np.dot(ray.direction, normal))),
                                     ray.color * surface_color), depth + 1)
            new_color += (add_color / RAY_NUMBER)
        return new_color
    # 灯光
    elif mode == 'light':
        return ray.color
    elif mode == 'back':
        return new_color


def raytrace():
    """
    光线跟踪主函数
    """
    ray = Ray([0, 0, 0], [1, 0, 0], [1, 1, 1])
    color = np.zeros(3, 'f')
    ray.position[:] = window.eye_position
    x_sample = [0.25, 0.25, -0.25, -0.25]
    y_sample = [0.25, -0.25, -0.25, 0.25]
    for y in range(window.height):
        process_bar(y / window.height, )
        for x in range(window.width):
            color[:] = 0
            for i in range(SAMPLE_NUMBER):
                ray.direction[:] = [x - window.width / 2 + x_sample[i],
                                    window.height / 2 - y + y_sample[i],
                                    -window.eye_position[2]]
                ray.normalize()
                color += traceray(ray, 0) / SAMPLE_NUMBER
                # 无交点
                window.color_map[y, x, :] = color
    window.show()


def raytrace_for_mp(x_range, y_range, mark):
    """
    用于多进程的光线追踪函数
    """
    result_array = np.zeros((int(length/2), length*2, 3), 'f')
    ray = Ray([0, 0, 0], [1, 0, 0], [1, 1, 1])
    color = np.zeros(3, 'f')
    ray.position[:] = window.eye_position
    # base_time用于估计剩余时间
    base_time = time.time()
    all_frame = length * length
    left_frame = all_frame
    done_frame = 0
    for y in y_range:
        for x in x_range:
            if mark == 3:
                frame_time = time.time() - base_time
                process_bar_left_time(done_frame / all_frame, frame_time * left_frame)
                done_frame += 1
                left_frame -= 1
                base_time = time.time()
            color[:] = 0
            for i in range(SAMPLE_NUMBER):
                ray.direction[:] = [x - window.width / 2 + random.random() - 0.5,
                                    window.height / 2 - y + random.random() - 0.5,
                                    -window.eye_position[2]]
                ray.normalize()
                color += traceray(ray, 0) / SAMPLE_NUMBER
                # print(window.color_map[y, x, :])
                result_array[y-int(mark*length/2), x, :] = color
    return result_array


def raytrace_mp():
    """
    多进程渲染
    """
    pool_number = 4
    pool = mp.Pool(pool_number)
    result = []
    for i in range(pool_number):
        # result_ = raytrace_for_mp(range(length*2), range(int(length/2*i), int(length/2*(i+1))), color_list[i], i)
        result.append(pool.apply_async(raytrace_for_mp, (range(length*2), range(int(length/2*i), int(length/2*(i+1))), i)))
    base_time = time.time()
    pool.close()
    pool.join()
    my_time = time.time() - base_time
    print('')
    print('total time:', my_time)
    j = 0
    # window.color_map[int(length/2*j):int(length/2*(j+1)), :, :] = result_
    for res in result:
        res_array = res.get()
        window.color_map[int(length/2*j):int(length/2*(j+1)), :, :] = res_array
        j += 1
    window.show()


def process_bar(value, left_time):
    bar_length = float(value * 100)
    h = str(int(left_time // 3600)) + ' h '
    m = str(int((left_time - h * 3600) // 60)) + ' m '
    s = str(int(left_time % 60)) + ' s '
    sys.stdout.write(str(bar_length) + ' left:' + h + m + s + '\r')
    sys.stdout.flush()


def process_bar_left_time(value, left_time):
    bar_length = float(value * 100)
    m = str(int(left_time // 60)) + ' m '
    s = str(left_time % 60)[:4] + ' s '
    sys.stdout.write(str(bar_length)[:5] + ' %' + ' left:' + m + s + '\r')
    sys.stdout.flush()


if __name__ == "__main__":
    # raytrace()
    raytrace_mp()
    # send_eamil('渲染完成')
